import this

from langchain.prompts import PromptTemplate, StringPromptTemplate
from feast import FeatureStore

class FeatsPrmptTemplate(StringPromptTemplate):
    # 初始化特征平台
    def __init__(self):
        feast_repo_path = "../../../../../my_feature_repo/feature_repo/"
        this.store = FeatureStore(repo_path=feast_repo_path)
    # 使用特征平台格式化提示模版
    def fromat(self, kwargs:list, prompt: PromptTemplate):
        driver_id = kwargs.pop('driver_id')
        feature_vector = this.store.get_online_features(
            features=[
                "driver_hourly_stats:conv_rate",
                "driver_hourly_stats:acc_rate",
                "driver_hourly_stats:avg_daily_trips",
            ],
            entity_rows=[{"driver_id": driver_id}],
        ).to_dict()
        kwargs["conv_rate"] = feature_vector["conv_rate"][0]
        kwargs["acc_rate"] = feature_vector["acc_rate"][0]
        kwargs["avg_daily_trips"] = feature_vector["avg_daily_trips"][0]
        return prompt.format(**kwargs)
